Skip to content

Add PointMass task and configurable optimizer for APG/PPO - #447

Open
yangchen73 wants to merge 7 commits into
mainfrom
yc/apg-env
Open

Add PointMass task and configurable optimizer for APG/PPO#447
yangchen73 wants to merge 7 commits into
mainfrom
yc/apg-env

Conversation

@yangchen73

Copy link
Copy Markdown
Collaborator

Description

  • Add a differentiable PointMass RL task, and train it with both APG and PPO through the same learning-env path
  • Make optimizer / LR scheduler configurable
  • Add a small script to plot PointMass trajectories

Type of change

  • New feature (non-breaking change which adds functionality)
  • Enhancement (non-breaking change which improves an existing functionality)
  • Breaking change (existing functionality will not work without user modification)

Screenshots

截屏2026-08-01 16 06 50 apg_best_ep3

Checklist

  • I have run the black . command to format the code base.
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • Dependencies have been updated, if applicable.

Copilot AI review requested due to automatic review settings August 1, 2026 20:33

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds a new lightweight, differentiable RL “PointMass” task and extends the RL stack so APG (differentiable rollouts) and PPO/GRPO (standard rollouts) can share a unified train-rl entrypoint via rollout-kind routing. It also introduces configurable optimizers/LR schedulers, unifies evaluation logic across trainers, and adds scripts/configs/docs to support benchmarking and visualization.

Changes:

  • Introduce PointMassRL as a registered lightweight learning environment and provide APG/PPO training configs plus a trajectory visualization script.
  • Add optimizer + LR-scheduler configuration plumbing through AlgorithmCfg/BaseAlgorithm, and update PPO/GRPO/APG to use the shared builders.
  • Unify deterministic evaluation and “best checkpoint” selection across Trainer and DifferentiableTrainer, with routing based on RolloutKind.

Reviewed changes

Copilot reviewed 51 out of 51 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
tests/learning/test_routing.py Adds coverage for routing trainer choice based on rollout_kind.
tests/learning/test_point_mass.py Adds extensive unit tests for PointMass env behavior + unified training entry.
tests/learning/test_optimizer.py Adds tests for optimizer and LR-scheduler builders/config coercion.
tests/learning/test_newton_planar_reach.py Updates APG config usage to the new nested OptimizerCfg.
tests/learning/test_evaluation.py Adds tests for unified deterministic evaluation semantics.
tests/learning/test_differentiable_trainer.py Extends checkpoint coverage to include LR scheduler state; adds close() to mocks.
tests/learning/test_differentiable_env.py Updates mock env protocol compliance with close().
tests/learning/test_apg.py Registers APG in algo registry, adds configurability + scheduler tests.
scripts/tutorials/rl/visualize_point_mass.py Adds a plotting script for deterministic PointMass trajectories from checkpoints.
scripts/benchmark/rl/tasks/point_mass.yaml Adds benchmark task definition for PointMass.
scripts/benchmark/rl/suites/point_mass.yaml Adds benchmark suite for PointMass across APG/PPO.
scripts/benchmark/rl/runtime.py Adds learning-env training path, unified evaluation usage, and routing integration.
scripts/benchmark/rl/algorithms/ppo.yaml Migrates PPO benchmark config to nested optimizer config.
scripts/benchmark/rl/algorithms/grpo.yaml Migrates GRPO benchmark config to nested optimizer config.
scripts/benchmark/rl/algorithms/apg.yaml Adds APG benchmark algorithm config.
embodichain/learning/rl/utils/trainer.py Unifies evaluation via evaluate_episodes, adds best-checkpoint selection and scheduler binding.
embodichain/learning/rl/utils/optimizer.py Introduces shared optimizer/LR-scheduler builders and registries.
embodichain/learning/rl/utils/helper.py Extends dict_to_tensordict to support tensor observations.
embodichain/learning/rl/utils/config.py Adds OptimizerCfg and LRSchedulerCfg into the shared algorithm config.
embodichain/learning/rl/utils/init.py Re-exports new optimizer/scheduler utilities and config objects.
embodichain/learning/rl/train.py Adds unified learning-env training path and rollout-kind based trainer selection.
embodichain/learning/rl/routing.py Adds get_trainer_class() to map rollout semantics to trainer type.
embodichain/learning/rl/models/init.py Extends MLP builder options and orthogonal initialization support.
embodichain/learning/rl/experimental/newton/train_planar_reach.py Migrates APG learning rate config to nested OptimizerCfg.
embodichain/learning/rl/evaluation.py Adds shared deterministic asynchronous evaluation helper for both trainers.
embodichain/learning/rl/env.py Adds learning-env registry (register_learning_env, build_learning_env) and shared protocol.
embodichain/learning/rl/differentiable_trainer.py Adds eval/best-checkpoint parity with Trainer, plus scheduler checkpointing.
embodichain/learning/rl/algo/ppo.py Switches PPO to shared optimization setup and scheduler stepping/logging.
embodichain/learning/rl/algo/grpo.py Switches GRPO to shared optimization setup and scheduler stepping/logging.
embodichain/learning/rl/algo/base.py Introduces RolloutKind + shared optimization/scheduler lifecycle hooks.
embodichain/learning/rl/algo/apg.py Switches APG to shared optimization setup, adds rollout kind and scheduler stepping/logging.
embodichain/learning/rl/algo/init.py Registers APG and normalizes nested optimizer/scheduler config mappings.
embodichain/learning/rl/init.py Re-exports new top-level RL APIs (learning env, evaluation, routing).
embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py Adds the differentiable PointMass environment and registers it as PointMassRL.
embodichain_tasks/embodichain_tasks/rl/basic/init.py Exports PointMass from the basic RL tasks package.
embodichain_tasks/configs/agents/rl/push_cube/train_config.json Migrates PPO config to nested optimizer block.
embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json Migrates GRPO config to nested optimizer block.
embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml Adds PPO training config for PointMass learning env.
embodichain_tasks/configs/agents/rl/basic/point_mass/train_apg.yaml Adds APG training config for PointMass learning env.
embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml Migrates PPO config to nested optimizer block.
embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json Migrates PPO config to nested optimizer block.
embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml Migrates GRPO config to nested optimizer block.
embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json Migrates GRPO config to nested optimizer block.
docs/source/tutorial/rl.rst Documents trainer.learning_env vs trainer.gym_config and adds PointMass examples.
docs/source/overview/rl/trainer.md Updates Trainer docs and adds DifferentiableTrainer parity notes.
docs/source/overview/rl/train_script.md Documents unified routing and PointMass usage in train.py.
docs/source/api_reference/embodichain/embodichain.learning.rl.rst Adds API reference sections for env/eval/routing + trainer.
docs/source/api_reference/embodichain/embodichain.learning.rl.algo.rst Updates algo overview to include APG and RolloutKind.
agent_context/topics/rl-learning/rl-learning.md Updates agent context docs for learning env registry, evaluation, and APG routing.
agent_context/MAP.yaml Updates “rl-learning” topic sources to include routing/evaluation/PointMass.
.agents/skills/add-task-env/SKILL.md Updates skill docs to reflect embodichain_tasks layout and learning-env tasks.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +121 to +127
metrics = info.get("metrics", {})
if isinstance(metrics, Mapping):
for name, value in metrics.items():
metric_values.setdefault(str(name), []).extend(
_selected_values(value, selected)
)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

Copilot AI review requested due to automatic review settings August 1, 2026 20:52

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 51 out of 51 changed files in this pull request and generated no new comments.

Copilot AI review requested due to automatic review settings August 1, 2026 22:31

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 52 out of 52 changed files in this pull request and generated no new comments.

Suppressed comments (4)

scripts/benchmark/rl/runtime.py:490

  • evaluate_checkpoint() binds an evaluation rollout buffer via env.set_rollout_buffer(), but the new evaluate_episodes() loop does not populate policy-side fields (notably previous actions) into that buffer. This can corrupt rewards/metrics for environments whose managers read prev actions from env.rollout_buffer (e.g. action_smoothness_penalty reads buffer["action"]). Prefer not binding a rollout buffer in this evaluation path unless the evaluator also writes the buffer fields each step.
    scripts/benchmark/rl/runtime.py:505
  • The on_step callback currently only compacts eval_rollout_buffer, but never writes action/reward/done into the buffer. After removing rollout-buffer binding for evaluation, this callback can be dropped and evaluate_episodes() called without on_step.
    embodichain/learning/rl/train.py:98
  • _build_learning_policy() computes policy_name but never uses it. This is dead code and can trigger lint/type-check warnings; remove it (or use it for validation) to keep the entry point clean.
    obs_dim = int(env.single_observation_space.shape[-1])
    action_dim = int(env.single_action_space.shape[-1])
    policy_name = policy_block["name"].lower()
    actor_cfg = policy_block.get("actor")
    critic_cfg = policy_block.get("critic")

embodichain_tasks/embodichain_tasks/rl/basic/point_mass.py:33

  • PointMassEnv is newly introduced, but it registers with override=True. This silently replaces any existing registration for "PointMassRL" and can mask accidental duplicate IDs during imports. Prefer the default (override=False) and rely on the registry's nested-package shadow detection for editable installs.

episode metrics.
- `save_checkpoint()`: Save model parameters and training state.

### DifferentiableTrainer

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should add docs for APG at algorithm section

@@ -0,0 +1,328 @@
# ----------------------------------------------------------------------------

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This script should not be included in official tutorial.

@yuecideng

Copy link
Copy Markdown
Contributor

Check that can DDP training work correctly for APG

@yuecideng

Copy link
Copy Markdown
Contributor

Also check whether the profiler can be correctly used for RL

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants